#!/usr/bin/env python

import os
from launch.substitutions import LaunchConfiguration
from launch.actions import DeclareLaunchArgument, ExecuteProcess
from launch.substitutions.path_join_substitution import PathJoinSubstitution
from launch import LaunchDescription
from launch_ros.actions import Node
import launch
from ament_index_python.packages import get_package_share_directory
from webots_ros2_driver.webots_launcher import WebotsLauncher
from webots_ros2_driver.webots_controller import WebotsController
from webots_ros2_driver.wait_for_controller_connection import WaitForControllerConnection


def generate_launch_description():
    # 获取包路径
    pkg_wheeltec = get_package_share_directory('turn_on_wheeltec_robot')
    pkg_wheeltec_urdf = get_package_share_directory('wheeltec_robot_urdf')

    # Launch配置
    world = LaunchConfiguration('world')
    mode = LaunchConfiguration('mode')
    use_sim_time = LaunchConfiguration('use_sim_time', default=True)

    # 启动Webots
    webots = WebotsLauncher(
        world=PathJoinSubstitution([pkg_wheeltec, 'worlds', world]),
        mode=mode,
        ros2_supervisor=True
    )

    # Robot State Publisher
    robot_description = os.path.join(pkg_wheeltec_urdf, 'urdf', 'my_tank_webots.urdf')
    
    robot_state_publisher = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        output='screen',
        parameters=[{
            'robot_description': open(robot_description, 'r').read(),
            'use_sim_time': use_sim_time,
            'publish_frequency': 50.0
        }],
    )

    # Base Footprint Publisher
    footprint_publisher = Node(
        package='tf2_ros',
        executable='static_transform_publisher',
        output='screen',
        arguments=['0', '0', '0', '0', '0', '0', 'base_link', 'base_footprint'],
    )

    # ROS Control Spawners
    controller_manager_timeout = ['--controller-manager-timeout', '100']
    controller_manager_prefix = 'python.exe' if os.name == 'nt' else ''
    
    diffdrive_controller_spawner = Node(
        package='controller_manager',
        executable='spawner',
        output='screen',
        prefix=controller_manager_prefix,
        arguments=['diffdrive_controller'] + controller_manager_timeout,
    )
    
    joint_state_broadcaster_spawner = Node(
        package='controller_manager',
        executable='spawner',
        output='screen',
        prefix=controller_manager_prefix,
        arguments=['joint_state_broadcaster'] + controller_manager_timeout,
    )
    
    ros_control_spawners = [diffdrive_controller_spawner, joint_state_broadcaster_spawner]

    # 配置文件路径
    ros2_control_params = os.path.join(pkg_wheeltec, 'config', 'ros2control.yml')

    # Tank Driver
    tank_driver = WebotsController(
        robot_name='tank',
        parameters=[
            {'robot_description': open(robot_description, 'r').read(),
             'use_sim_time': use_sim_time,
             'set_robot_state_publisher': True},
            ros2_control_params
        ],
        remappings=[
            ('/diffdrive_controller/cmd_vel_unstamped', '/cmd_vel'),
            ('/diffdrive_controller/odom', '/odom')
        ],
        output='screen'
    )

    # 等待控制器连接
    waiting_nodes = WaitForControllerConnection(
        target_driver=tank_driver,
        nodes_to_start=ros_control_spawners
    )

    return LaunchDescription([
        # 声明参数
        DeclareLaunchArgument(
            'world',
            default_value='factory.wbt',
            description='Choose one of the world files from worlds directory'
        ),
        DeclareLaunchArgument(
            'mode',
            default_value='realtime',
            description='Webots startup mode'
        ),
        
        # 添加延迟节点
        ExecuteProcess(
            cmd=['sleep', '5'],
            name='delay_webots_start'
        ),
        
        # 添加节点
        webots,
        webots._supervisor,
        
        # 等待Webots启动
        ExecuteProcess(
            cmd=['sleep', '2'],
            name='delay_controller_start'
        ),
        
        robot_state_publisher,
        footprint_publisher,
        tank_driver,
        waiting_nodes,

        # 关闭处理
        launch.actions.RegisterEventHandler(
            event_handler=launch.event_handlers.OnProcessExit(
                target_action=webots,
                on_exit=[launch.actions.EmitEvent(event=launch.events.Shutdown())],
            )
        ),
    ])
